from flask import Blueprint, jsonify, request
from app.models.knowledge_base import KnowledgeBase
from loguru import logger


kb_bp = Blueprint('kb', __name__, url_prefix='/api/kb')
def init_kb_api(kb: KnowledgeBase):
    """初始化知识库API"""
    
    @kb_bp.route('/build', methods=['POST'])
    def build_knowledge_base():
        logger.info("############ 构建知识库... ############")
        """构建知识库"""
        try:
            kb.build_vector_store()
            return jsonify({"status": "success", "message": "知识库构建成功"}), 200
            logger.info("############ 构建知识库完成 ############")
        except Exception as e:
            logger.error(f"构建知识库失败: {e}")
            return jsonify({"status": "error", "message": str(e)}), 500
    
    @kb_bp.route('/query', methods=['POST'])
    def query_knowledge_base():
        logger.info("############ 查询知识库... ############")
        """查询知识库"""
        data = request.get_json()
        query = data.get('query')
        k = data.get('k', 3)
        
        if not query:
            return jsonify({"status": "error", "message": "查询内容不能为空"}), 400
        logger.info(f"##############知识库的参数：{query} ############")
        try:
            results = kb.query(query, k)
            logger.info(f"############## 知识库查询完成！结果：{results} ############")
            return jsonify({"status": "success", "results": results}), 200
        except Exception as e:
            logger.error(f"查询知识库失败: {e}")
            return jsonify({"status": "error", "message": str(e)}), 500
    return kb_bp